import numpy as np
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import random

# 尝试导入matplotlib，如果失败则跳过绘图功能
try:
    import matplotlib.pyplot as plt
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    print("警告: matplotlib未安装，将跳过绘图功能")


def generate_sine_test_data(filename="sine_test_data.csv", num_points=1000, noise_level=0.1, add_label=False, label_categories=None, label_distribution=None):
    """
    生成测试数据，包含时间戳和三列正弦波叠加高斯噪声的数据

    参数:
    filename: 输出文件名
    num_points: 数据点数量
    noise_level: 噪声水平（标准差）
    add_label: 是否添加标签列（可选）
    label_categories: 标签类别列表，默认为['A', 'B', 'C']
    label_distribution: 标签分布比例，默认为均匀分布
    """
    # 生成时间戳（从当前时间开始，每隔1分钟一个数据点）
    start_time = datetime.now()
    timestamps = [start_time + timedelta(minutes=i) for i in range(num_points)]
    unix_timestamps = [int(ts.timestamp()) for ts in timestamps]

    # 生成三列正弦波数据，每列使用不同的频率和相位
    x = np.linspace(0, 4 * np.pi, num_points)

    # 第二列数据：基本频率的正弦波
    sine_1 = np.sin(x) + np.random.normal(0, noise_level, num_points)

    # 第三列数据：两倍频率的正弦波
    sine_2 = np.sin(2 * x + np.pi / 4) + np.random.normal(0, noise_level, num_points)

    # 第四列数据：三倍频率的正弦波
    sine_3 = np.sin(3 * x + np.pi / 2) + np.random.normal(0, noise_level, num_points)

    # 创建基础数据字典
    data_dict = {
        'timestamp': unix_timestamps,
        'sensor_1': sine_1,
        'sensor_2': sine_2,
        'sensor_3': sine_3
    }
    
    # 如果添加标签列
    if add_label:
        # 设置默认标签类别
        if label_categories is None:
            label_categories = ['A', 'B', 'C']
        
        # 设置默认分布（均匀分布）
        if label_distribution is None:
            label_distribution = [1.0 / len(label_categories)] * len(label_categories)
        
        # 生成标签数据
        labels = np.random.choice(label_categories, size=num_points, p=label_distribution)
        data_dict['label'] = labels
    
    # 创建DataFrame
    df = pd.DataFrame(data_dict)

    # 保存到CSV文件
    df.to_csv(filename, index=False)
    print(f"数据已保存到 {filename}")

    # 绘制数据图表（仅在matplotlib可用时）
    if MATPLOTLIB_AVAILABLE:
        fig, axes = plt.subplots(3, 1, figsize=(12, 8))

        axes[0].plot(unix_timestamps, sine_1, label='Sensor 1')
        axes[0].set_title('Sensor 1 Data (sin(x))')
        axes[0].set_ylabel('Value')
        axes[0].grid(True)

        axes[1].plot(unix_timestamps, sine_2, label='Sensor 2', color='orange')
        axes[1].set_title('Sensor 2 Data (sin(2x + π/4))')
        axes[1].set_ylabel('Value')
        axes[1].grid(True)

        axes[2].plot(unix_timestamps, sine_3, label='Sensor 3', color='green')
        axes[2].set_title('Sensor 3 Data (sin(3x + π/2))')
        axes[2].set_xlabel('Timestamp')
        axes[2].set_ylabel('Value')
        axes[2].grid(True)

        plt.tight_layout()
        plt.savefig('sine_test_data_plot.png')
        plt.show()
    else:
        print("matplotlib不可用，跳过绘图")

    return df


def generate_label_test_data(filename="label_test_data.csv", num_points=1000, label_columns=None):
    """
    生成纯标签测试数据，用于测试类别统计插件
    
    参数:
    filename: 输出文件名
    num_points: 数据点数量
    label_columns: 标签列配置，例如：
        {
            'category_1': {'categories': ['A', 'B', 'C'], 'distribution': [0.4, 0.3, 0.3]},
            'category_2': {'categories': ['X', 'Y'], 'distribution': [0.6, 0.4]}
        }
    """
    if label_columns is None:
        # 默认配置，对应data_pipeline.yaml中的配置
        label_columns = {
            'category_1': {'categories': ['A', 'B', 'C'], 'distribution': [0.4, 0.3, 0.3]},
            'category_2': {'categories': ['X', 'Y'], 'distribution': [0.6, 0.4]}
        }
    
    # 生成时间戳
    start_time = datetime.now()
    timestamps = [start_time + timedelta(minutes=i) for i in range(num_points)]
    unix_timestamps = [int(ts.timestamp()) for ts in timestamps]
    
    # 创建数据字典
    data_dict = {'timestamp': unix_timestamps}
    
    # 生成每列标签数据
    for col_name, col_config in label_columns.items():
        categories = col_config['categories']
        distribution = col_config['distribution']
        
        # 生成标签数据
        labels = np.random.choice(categories, size=num_points, p=distribution)
        data_dict[col_name] = labels
    
    # 创建DataFrame
    df = pd.DataFrame(data_dict)
    
    # 保存到CSV文件
    df.to_csv(filename, index=False)
    print(f"标签测试数据已保存到 {filename}")
    
    # 显示统计信息
    print("\n标签统计:")
    for col_name in label_columns.keys():
        print(f"\n{col_name}:")
        print(df[col_name].value_counts().sort_index())
    
    return df


def generate_component_sensor_data(filename="component_sensor_data.csv", num_points=1000, 
                                  component_name="machine_001", sampling_interval=1, 
                                  sensor1_noise_level=0.1, sensor2_noise_level=0.1):
    """
    生成一个部件的两个传感器采集数据集，包含时间戳、部件名和传感器采集值

    参数:
    filename: 输出文件名
    num_points: 数据点数量
    component_name: 部件名称
    sampling_interval: 采样间隔（秒）
    sensor1_noise_level: 第一个传感器的噪声水平
    sensor2_noise_level: 第二个传感器的噪声水平
    """
    # 生成时间戳（从当前时间开始，按指定间隔采样）
    start_time = datetime.now()
    timestamps = [start_time + timedelta(seconds=i*sampling_interval) for i in range(num_points)]
    unix_timestamps = [int(ts.timestamp()) for ts in timestamps]
    human_timestamps = [ts.strftime('%Y-%m-%d %H:%M:%S') for ts in timestamps]

    # 生成传感器1的数据（基于正弦波，模拟正常运行状态）
    x = np.linspace(0, 10 * np.pi, num_points)
    base_value1 = 50 + 10 * np.sin(x)
    noise1 = np.random.normal(0, sensor1_noise_level, num_points)
    sensor1_values = base_value1 + noise1

    # 生成传感器2的数据（基于正弦波，但有不同的相位和振幅，模拟与传感器1相关但不完全相同的数据）
    base_value2 = 60 + 8 * np.sin(x + np.pi/3)
    noise2 = np.random.normal(0, sensor2_noise_level, num_points)
    sensor2_values = base_value2 + noise2

    # 偶尔添加一些异常值（模拟传感器故障或异常情况）
    anomaly_indices = random.sample(range(num_points), int(num_points * 0.02))  # 2%的数据点作为异常
    for idx in anomaly_indices:
        if random.random() < 0.5:
            # 传感器1出现异常
            sensor1_values[idx] = sensor1_values[idx] * (0.5 + random.random())  # 随机增加或减少值
        else:
            # 传感器2出现异常
            sensor2_values[idx] = sensor2_values[idx] * (0.5 + random.random())  # 随机增加或减少值

    # 创建数据字典
    data_dict = {
        'timestamp_unix': unix_timestamps,
        'timestamp_human': human_timestamps,
        'component_name': [component_name] * num_points,
        'sensor_1': sensor1_values,
        'sensor_2': sensor2_values
    }

    # 创建DataFrame
    df = pd.DataFrame(data_dict)

    # 保存到CSV文件
    df.to_csv(filename, index=False)
    print(f"部件传感器数据已保存到 {filename}")

    # 绘制数据图表（仅在matplotlib可用时）
    if MATPLOTLIB_AVAILABLE:
        fig, axes = plt.subplots(2, 1, figsize=(12, 8))

        axes[0].plot(human_timestamps[::10], sensor1_values[::10], label='Sensor 1', marker='o')  # 每10个点显示一个标记
        axes[0].set_title(f'Component: {component_name} - Sensor 1 Data')
        axes[0].set_ylabel('Value')
        axes[0].grid(True)
        # 只显示部分时间戳标签，避免重叠
        for label in axes[0].get_xticklabels()[::2]:
            label.set_visible(False)

        axes[1].plot(human_timestamps[::10], sensor2_values[::10], label='Sensor 2', color='orange', marker='s')
        axes[1].set_title(f'Component: {component_name} - Sensor 2 Data')
        axes[1].set_xlabel('Timestamp')
        axes[1].set_ylabel('Value')
        axes[1].grid(True)
        # 只显示部分时间戳标签，避免重叠
        for label in axes[1].get_xticklabels()[::2]:
            label.set_visible(False)

        plt.tight_layout()
        plt.savefig('component_sensor_data_plot.png')
        plt.show()
    else:
        print("matplotlib不可用，跳过绘图")

    # 显示数据统计信息
    print(f"\n数据统计信息:")
    print(f"总数据点数量: {num_points}")
    print(f"采样时间范围: {human_timestamps[0]} 到 {human_timestamps[-1]}")
    print(f"传感器1数据范围: {min(sensor1_values):.2f} - {max(sensor1_values):.2f}")
    print(f"传感器1平均值: {np.mean(sensor1_values):.2f}")
    print(f"传感器1标准差: {np.std(sensor1_values):.2f}")
    print(f"传感器2数据范围: {min(sensor2_values):.2f} - {max(sensor2_values):.2f}")
    print(f"传感器2平均值: {np.mean(sensor2_values):.2f}")
    print(f"传感器2标准差: {np.std(sensor2_values):.2f}")
    print(f"异常数据点数量: {len(anomaly_indices)} ({len(anomaly_indices)/num_points*100:.1f}%)")

    # 显示前几行数据示例
    print(f"\n数据前5行:")
    print(df.head())

    return df

def generate_component_sensor_datasets(file1="component1_sensor_data.csv", file2="component2_sensor_data.csv", 
                                      num_points=1000, sampling_interval=1, 
                                      component1_name="machine_001", component2_name="machine_001"):
    """
    生成两个不同部件的传感器采集数据集，每个数据集包含时间戳、部件名和传感器采集值

    参数:
    file1: 第一个部件数据集文件名
    file2: 第二个部件数据集文件名
    num_points: 每个数据集的数据点数量
    sampling_interval: 采样间隔（秒）
    component1_name: 第一个部件的名称
    component2_name: 第二个部件的名称
    """
    # 生成时间戳（从当前时间开始，按指定间隔采样）
    start_time = datetime.now()
    timestamps = [start_time + timedelta(seconds=i*sampling_interval) for i in range(num_points)]
    unix_timestamps = [int(ts.timestamp()) for ts in timestamps]
    human_timestamps = [ts.strftime('%Y-%m-%d %H:%M:%S') for ts in timestamps]

    # ============== 第一个部件的数据集 ==============
    # 生成第一个部件的传感器数据（基于正弦波，模拟正常运行状态）
    x = np.linspace(0, 10 * np.pi, num_points)
    
    # 传感器1数据（基础值 + 噪声）
    base_value1_comp1 = 50 + 10 * np.sin(x)
    noise1_comp1 = np.random.normal(0, 0.1, num_points)
    sensor1_comp1 = base_value1_comp1 + noise1_comp1
    
    # 偶尔添加一些异常值（模拟传感器故障或异常情况）
    anomaly_indices1 = random.sample(range(num_points), int(num_points * 0.02))  # 2%的数据点作为异常
    for idx in anomaly_indices1:
        # 随机增加或减少值
        sensor1_comp1[idx] = sensor1_comp1[idx] * (0.5 + random.random())
    
    # 创建第一个数据集的字典
    data_dict1 = {
        'timestamp_unix': unix_timestamps,
        'timestamp_human': human_timestamps,
        'component_name': [component1_name] * num_points,
        'sensor_value_1': sensor1_comp1
    }
    
    # ============== 第二个部件的数据集 ==============
    # 生成第二个部件的传感器数据（与第一个部件有相关性但不完全相同）
    # 使用不同的基础值、振幅和相位
    base_value_comp2 = 45 + 12 * np.sin(x + np.pi/4)
    noise_comp2 = np.random.normal(0, 0.12, num_points)
    sensor_comp2 = base_value_comp2 + noise_comp2

    # 偶尔添加一些异常值
    anomaly_indices2 = random.sample(range(num_points), int(num_points * 0.02))  # 2%的数据点作为异常
    for idx in anomaly_indices2:
        # 随机增加或减少值
        sensor_comp2[idx] = sensor_comp2[idx] * (0.5 + random.random())

    # 创建第二个数据集的字典
    data_dict2 = {
        'timestamp_unix': unix_timestamps,
        'timestamp_human': human_timestamps,
        'component_name': [component2_name] * num_points,
        'sensor_value_2': sensor_comp2
    }
    
    # 创建DataFrame
    df1 = pd.DataFrame(data_dict1)
    df2 = pd.DataFrame(data_dict2)
    
    # 保存到CSV文件
    df1.to_csv(file1, index=False)
    df2.to_csv(file2, index=False)
    
    print(f"\n=== 生成部件传感器测试数据集 ===")
    print(f"第一个部件数据集已保存到 {file1} (共{len(df1)}条记录)")
    print(f"第二个部件数据集已保存到 {file2} (共{len(df2)}条记录)")
    
    # 显示第一个数据集的统计信息
    print(f"\n{component1_name} 数据统计:")
    print(f"数据范围: {min(sensor1_comp1):.2f} - {max(sensor1_comp1):.2f}")
    print(f"平均值: {np.mean(sensor1_comp1):.2f}")
    print(f"标准差: {np.std(sensor1_comp1):.2f}")
    print(f"异常数据点数量: {len(anomaly_indices1)} ({len(anomaly_indices1)/num_points*100:.1f}%)")
    
    # 显示第二个数据集的统计信息
    print(f"\n{component2_name} 数据统计:")
    print(f"数据范围: {min(sensor_comp2):.2f} - {max(sensor_comp2):.2f}")
    print(f"平均值: {np.mean(sensor_comp2):.2f}")
    print(f"标准差: {np.std(sensor_comp2):.2f}")
    print(f"异常数据点数量: {len(anomaly_indices2)} ({len(anomaly_indices2)/num_points*100:.1f}%)")
    
    # 显示数据集前几行示例
    print(f"\n{file1} 前5行:")
    print(df1.head())
    print(f"\n{file2} 前5行:")
    print(df2.head())
    
    return df1, df2


if __name__ == "__main__":
    # 生成带标签的测试数据
    print("=== 生成带标签的测试数据 ===")
    data_with_label = generate_sine_test_data(
        "sine_test_data_with_label.csv", 
        num_points=1000, 
        noise_level=0.1,
        add_label=True,
        label_categories=['A', 'B', 'C'],
        label_distribution=[0.4, 0.3, 0.3]  # A类40%，B类30%，C类30%
    )
    print("数据前10行:")
    print(data_with_label.head(10))
    print(f"\n标签统计:")
    print(data_with_label['label'].value_counts())
    
    print("\n=== 生成不带标签的测试数据 ===")
    # 生成不带标签的测试数据
    data_without_label = generate_sine_test_data("sine_test_data.csv", num_points=1000, noise_level=0.1)
    print("数据前10行:")
    print(data_without_label.head(10))
    
    print("\n=== 生成纯标签测试数据 ===")
    # 生成纯标签测试数据，对应data_pipeline.yaml配置
    label_data = generate_label_test_data(
        "label_test_data.csv", 
        num_points=1000,
        label_columns={
            'category_1': {'categories': ['A', 'B', 'C'], 'distribution': [0.4, 0.3, 0.3]},
            'category_2': {'categories': ['X', 'Y'], 'distribution': [0.6, 0.4]}
        }
    )
    print("\n数据前10行:")
    print(label_data.head(10))
    
    # 生成用于数据合并的测试数据
    generate_component_sensor_datasets()
    
    print("\n=== 生成部件传感器测试数据 ===")
    # 生成一个部件的两个传感器采集数据集
    component_data = generate_component_sensor_data(
        "component_sensor_data.csv",
        num_points=1000,
        component_name="machine_001",
        sampling_interval=1,
        sensor1_noise_level=0.1,
        sensor2_noise_level=0.1
    )